import tensorflow
from keras.layers import Input, Dense
from keras.models import Sequential, Model
import numpy as np
import pandas as pd
import random
from sklearn.preprocessing import minmax_scale


def read_label_data(DataFrame):
    feature_1 = []
    feature_0 = []
    label_1 = []
    label_0 = []

    for row in range(len(DataFrame)):
        temp_feature = []
        if DataFrame[DataFrame.columns[len(DataFrame.columns) - 1]][row] == 1:
            for col in range(len(DataFrame.columns) - 1):
                temp_feature.append(DataFrame[DataFrame.columns[col]][row])
            feature_1.append(temp_feature)
            label_1.append(1)
        else:
            for col in range(len(DataFrame.columns) - 1):
                temp_feature.append(DataFrame[DataFrame.columns[col]][row])
            feature_0.append(temp_feature)
            label_0.append(0)

    return feature_1, feature_0, label_1, label_0


def get_list(list_list, index):
    list_list = list_list
    index = index
    result_list = []
    for x in index:
        result_list.append(list_list[x])
    return result_list


data = pd.read_csv("australian.dat", sep=" ", header=None)

feature_bad, feature_good, label_bad, label_good = read_label_data(data)

# 归一化
feature_good = minmax_scale(feature_good)
feature_bad = minmax_scale(feature_bad)


class GAN(object):
    def __init__(self):
        # 特征维度
        self.feature_shape = len(data.columns) - 1
        # 优化器
        optimizer = tensorflow.keras.optimizers.SGD()
        # 定义D，设置训练参数
        self.discriminator = self.build_discriminator()
        self.discriminator.compile(loss='BinaryFocalCrossentropy',
                                   optimizer=optimizer,
                                   metrics=['accuracy']
                                   )
        # 定义G
        self.generator = self.build_generator()
        # 构建GAN结构
        gan_input = Input(shape=(self.feature_shape, ))
        generated_feature = self.generator(gan_input)
        self.discriminator.trainable = False
        validity = self.discriminator(generated_feature)
        self.gan = Model(gan_input, validity)
        self.gan.compile(loss='BinaryFocalCrossentropy',
                         optimizer=optimizer
                         )

    def build_generator(self):
        model = Sequential()
        model.add(Dense(self.feature_shape, activation="LeakyReLU"))
        model.add(Dense(32, activation="LeakyReLU"))
        model.add(Dense(64, activation="LeakyReLU"))
        model.add(Dense(128, activation="LeakyReLU"))
        model.add(Dense(64, activation="LeakyReLU"))
        model.add(Dense(32, activation="LeakyReLU"))
        model.add(Dense(self.feature_shape))

        noise = Input(shape=(self.feature_shape, ))
        generated_feature = model(noise)

        return Model(noise, generated_feature)

    def build_discriminator(self):
        model = Sequential()
        model.add(Dense(self.feature_shape, activation="LeakyReLU"))
        model.add(Dense(int(self.feature_shape/2), activation="LeakyReLU"))
        model.add(Dense(1, activation="sigmoid"))

        feature = Input(shape=(self.feature_shape,))
        labels = model(feature)

        return Model(feature, labels)

    def train_GAN(self, epochs=30000, batch_size=32, sample_interval=10):

        for epoch in range(epochs):
            real_feature_index = random.sample(range(0, len(feature_bad)), batch_size)
            real_feature = get_list(feature_bad, real_feature_index)
            real_label = get_list(label_bad, real_feature_index)

            mix_feature = real_feature
            mix_label = real_label
            for x in range(batch_size):
                latent_space_sample = [real_feature[random.sample(range(0, batch_size), 1)[0]]]
                # 用G给原始样本添加扰动
                mix_feature.append(self.generator.predict(np.array(latent_space_sample))[0].tolist())
                mix_label.append(int(np.zeros(1)))

            feature_batch = mix_feature
            label_batch = mix_label

            # 训练D：把真的1样本预测为1，把生成的1样本预测为0
            self.discriminator.train_on_batch(np.array(feature_batch), np.array(label_batch))

            GAN_feature = []
            GAN_label = []
            for x in range(batch_size):
                latent_space_sample = [real_feature[random.sample(range(0, batch_size), 1)[0]]]
                # 用G给原始样本添加扰动
                GAN_feature.append(self.generator.predict(np.array(latent_space_sample))[0].tolist())
                GAN_label.append(int(np.ones(1)))

            # 训练G：让GAN把生成的1样本预测为1
            self.gan.train_on_batch(np.array(GAN_feature), np.array(GAN_label))

            if epoch % sample_interval == 0:
                latent_space_sample = [real_feature[random.sample(range(0, batch_size), 1)[0]]]
                GAN_result = self.generator.predict(np.array(latent_space_sample))
                # 把生成的1样本标记为0输出
                G_sample = np.append(GAN_result[0], 0)
                print(G_sample)
                new_data = pd.DataFrame([G_sample])
                new_data.to_csv('G_data.csv',
                                mode='a',
                                header=False,
                                float_format="%.5f",
                                index=False
                                )

            print("epochs:{}".format(epoch))


if __name__ == '__main__':
    gan = GAN()
    gan.train_GAN(epochs=30000, batch_size=32, sample_interval=10)
